import rqdatac as rq
import pandas as pd
from datetime import timedelta

night_secs = ["CU", "AL", "AG", "SP", "RU", "SA", "FG", "I",
              "RB", "J", "BU", "TA", "MA", "V", "P", "M", "C",
              "CF", "SR", "A"]
day_secs = []
SECS = night_secs + day_secs
#SECS = ["RB", "SM"]
day_sec_num = len(SECS)
night_secs_num = len(night_secs)
start_date = "20220701"
end_date = "20230225"

rq.init("13570866213", "39314656")
secs = [f"{sec}99" for sec in SECS]
#n_secs = [f"{sec}99" for sec in night_secs]
data_df = rq.get_price(secs, frequency="1m", fields=["close"], start_date=start_date, end_date=end_date, adjust_type="none")
data_df.reset_index(drop=False, inplace=True)
data_df["datetime"] = data_df["datetime"] - timedelta(minutes=1)
data_df["date"] = data_df["datetime"].dt.date
data_df["hour"] = data_df["datetime"].dt.hour
data_df.set_index("datetime", inplace=True)

data_df = data_df.loc[(data_df["hour"] >= 9) & (data_df["hour"] <= 22)]
dates = data_df['date'].drop_duplicates(keep="first")
#data_df.drop(["hour", "date"], axis=1, inplace=True)
sec_datas = {}
for sec in secs:
    sec_data = data_df.loc[data_df["order_book_id"] == sec, "close"]
    sec_data.dropna(inplace=True)
    sec_data_1 = sec_data.shift(1)
    d = sec_data - sec_data_1
    sec_data = d / sec_data_1
    sec_datas[sec] = sec_data

data_df = pd.DataFrame(sec_datas)
data_df.reset_index(drop=False, inplace=True)
data_df["date"] = data_df["datetime"].dt.date
data_df["hour"] = data_df["datetime"].dt.hour
data_df.set_index("datetime", inplace=True)
night_data = data_df.loc[(data_df["hour"] >= 21) & (data_df["hour"] <= 22)]
day_data = data_df.loc[(data_df["hour"] >= 9) & (data_df["hour"] <= 15)]


def calculate_corr(ori_df):
    if ori_df is None:
        return None
    if ori_df.empty:
        return None
    ori_df.fillna(method="bfill", inplace=True)
    ori_df.fillna(method="ffill", inplace=True)
    ori_df.dropna(axis=1, inplace=True)
    sec_num = ori_df.shape[1]
    if sec_num <= 1:
        return None
    #df = ori_df[secs]
    df_corr = ori_df.corr()
    df_corr["corr"] = (df_corr.sum(axis=1) - 1) / (sec_num - 1)
    return df_corr["corr"].mean()


results = []
for date in dates:
    night_df = night_data.loc[night_data["date"] == date, secs]
    night_corr = calculate_corr(night_df)
    day_df = day_data.loc[day_data["date"] == date, secs]
    day_df = day_df.iloc[1:]
    day_corr = calculate_corr(day_df)
    results.append({"date": date, "day": day_corr, "night": night_corr})

result_df = pd.DataFrame(results)
result_df.to_csv("E:\\corr_result_drop1_20230225.csv")
